#include "unifiedMathFunctions.hpp"
#include "reconstruction.hpp"
#include "bi_reconstruction.hpp"
#include "hllc.hpp"

#ifdef BLOCKWISE_INTEGRATED_PRECALC_DIFFS
#include "bi_precalc_diffs.hpp"
#endif

namespace hms::swe
{


template<typename F, typename Z, typename D, typename N>
void slopeSourceToFlux(
	F& flux,
	EdgesValues& ev,
	const Z& zCellLR,
	const D& dCellLR,
	const N& normals
){
	for (int i{0}; i<2; ++i){
		ev.zLR[i] = min( ev.hLR[i], ev.zLR[i] );

		/* the equation is f_S,x = -nx * 0.5*g * (dL + dCellL) * (zL - zCellL)
		 * and f_S,y = -ny * 0.5*g * (dL + ...) * ...
		 * The negative sign is added later in registerFlux.
		 */
		flux[i].template bottomRows<2>() +=
			normals.array() * (phys::g / 2) *
			( (ev.dLR[i] + dCellLR[i]) * (ev.zLR[i] - zCellLR[i]) )
			.template replicate<2,1>();
	}
}


template<typename MeshType>
template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
// HMS_INLINE
EdgesValues SolverThread<MeshType>::reconstruction(
	const EdgeRow<recType,orien>& er,
	const MeshBase<MeshType>& mesh,
	const Block<loc>& block,
	const ArrayXXsMap& qBlock,
	const Ref<const ArrayXXs>& dBlock,
	const Ref<const ArrayXXs>& zBlock,
	const Ref<const ArrayXXs>& hBlock,
	const Ref<const ArrayXXs>& vBlock
){
	using RT = ReconstructionType;
	using ER = EdgeRange;
	using S = SolverType;

	EdgesValues ev { mapEdgesValues( m_edgesValsBuf, er.nEdges ) };

	Pair<Array1XsCMap>
		dCellLR { block.template mapLR<Array1XsCMap,false>(dBlock, er) },
		zCellLR { block.template mapLR<Array1XsCMap,false>(zBlock, er) };

	Pair<Array2XsCMap>
		vCellLR { block.template mapLR<Array2XsCMap,false>(vBlock, er) };

	if constexpr ( isOrtho && recType != RT::firstOrder ){
		Pair<Array1XsCMap>
		hCellLR { block.template mapLR<Array1XsCMap,false>(hBlock, er) };

		/* because q has three rows, but we only want to map to the bottom two,
		 * we must specify the stride to use. The inner stride
		 * (distance between entries of one column) remains at 0
		 * (natural stride, I think that has the same effect as 1), but the
		 * outer stride, which is the pointer distance between columns, is the
		 * number of entries per column in the memory to be mapped, which is
		 * 3. */
		using QMapBottom = Eigen::Map< const Array2Xs,0,Eigen::OuterStride<3> >;

		Pair<QMapBottom> qCellLR {
			block.template mapLR<QMapBottom,Array3XsMap,1,false>(qBlock, er)
		};

		if constexpr (recType == RT::secondOrder){
			Pair<Array1XsCMap>
			dCellLLRR { block.template mapLR<Array1XsCMap,true>(dBlock, er) },
			hCellLLRR { block.template mapLR<Array1XsCMap,true>(hBlock, er) };

			Pair<QMapBottom> qCellLLRR {
				block.template mapLR<QMapBottom,Array3XsMap,1,true>(qBlock, er)
			};

			/* move this to RectCells */
			auto dx { [&](){
				if constexpr ( isUniMesh ){
					return Array1s{};
				} else if constexpr ( isRectMesh ){
					if constexpr ( orien == ER::horizontal ){
						/* NEEDS TESTING */
						return mesh.cells().dy().segment(
							block.yBeg() - block.cb + er.i, 3
						);
					} else {
						/* NEEDS TESTING */
						return mesh.cells().dx().segment(
							block.xBeg() - block.cl, er.nEdges + 2
						);
					}
				}
			}() };
			auto
				dxL { dxQuotient<MeshType, orien, 0>(dx) },
				dxR { dxQuotient<MeshType, orien, 1>(dx) };

			/* precalculate differences */
			#ifdef BLOCKWISE_INTEGRATED_PRECALC_DIFFS
			auto [dDiffL, dDiffM, dDiffR] =
				phiDiff<loc,1>( m_phiDiffBuf, 0, dCellLR, dCellLLRR, er );
			auto [hDiffL, hDiffM, hDiffR] =
				phiDiff<loc,1>( m_phiDiffBuf, 3, hCellLR, hCellLLRR, er );
			auto [qDiffL, qDiffM, qDiffR] =
				phiDiff<loc,2>( m_phiDiffBuf, 6, qCellLR, qCellLLRR, er );
			assert( dDiffR.data() + dDiffR.size() <= hDiffL.data() );
			assert( hDiffR.data() + hDiffR.size() <= qDiffL.data() );

			reconstruction::highOrder<S>(
				ev,
				dCellLR, dDiffL, dDiffM, dDiffR,
				hCellLR, hDiffL, hDiffM, hDiffR,
				qCellLR, qDiffL, qDiffM, qDiffR,
				dxL, dxR, solver().settings().limiterFunction
			);

			#else

			reconstruction::highOrder<S>(
				ev,
				dCellLR, dCellLLRR,
				hCellLR, hCellLLRR,
				qCellLR, qCellLLRR,
				dxL, dxR,
				solver().settings().limiterFunction
			);
			#endif
		} else if constexpr (recType == RT::central) {
			reconstruction::central<S>(ev, dCellLR, hCellLR, qCellLR);
		}
		/* murillo's condition is checked for both the central method as well
		 * as muscl */
		murilloCheck<S>(ev, dCellLR, zCellLR, vCellLR);
	} else {
		reconstruction::firstOrder<S>( ev, dCellLR, zCellLR, vCellLR );
	}
	reconstruction::hydrostatic_firstOrder(ev);

	return ev;
}

template<typename MeshType>
template<EdgeRange orien>
Pair<Array3XsMap> SolverThread<MeshType>::hllcScheme(
	const EdgesValues& ev,
	Index nEdges
){
	using ER = EdgeRange;
	using S = SolverType;

	/* vn is the velocity component normal to the edge,
	 * vp is the velocity component parallel to the edge. */
	auto vn = [&](int side) -> decltype(auto) {
		if constexpr (isOrtho){
			if constexpr ( orien == ER::horizontal )
				return ev.vLR[side].row(1);
			else
				return ev.vLR[side].row(0);
		} else {
			assert(false);
		}
	};
	auto vp = [&](int side) -> decltype(auto) {
		if constexpr (isOrtho){
			if constexpr ( orien == ER::horizontal )
				return - ev.vLR[side].row(0);
			else
				return ev.vLR[side].row(1);
		}
	};
	auto vnL { vn(left) };
	auto vnR { vn(right) };
	auto vpL { vp(left) };
	auto vpR { vp(right) };

	const Array1XsMap
		& dL {ev.dLR[left ]},
		& dR {ev.dLR[right]};

	assert( m_hllcBuf.size() >= nEdges * nHllcComp );

	const auto& normal { rectNormals<orien>(nEdges) };

	assert( m_fluxLineBuf.cols() >= 2 );
	assert( m_fluxLineBuf.rows() >= 3*nEdges );

	Pair<Array3XsMap> flux {{
		{ m_fluxLineBuf.col(left ).data(), 3, nEdges },
		{ m_fluxLineBuf.col(right).data(), 3, nEdges }
	}};

	hllc::hllcScheme<S, orien>( flux, m_hllcBuf, dL, dR, vnL, vnR, vpL, vpR, normal );

	return flux;
}

template<typename MeshType>
template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
void SolverThread<MeshType>::fluxScheme_vec(
	EdgeRow<recType,orien>&& er,
	const Array2i& lrBeg_block,
	const MeshBase<MeshType>& mesh,
	const Block<loc>& block,
	[[maybe_unused]] const Array3XsMap& q,
	[[maybe_unused]] const ArrayXXsMap& qBlock,
	Array3XsMap& fluxes,
	const Ref<const ArrayXXs>& dBlock,
	const Ref<const ArrayXXs>& zBlock,
	const Ref<const ArrayXXs>& hBlock,
	const Ref<const ArrayXXs>& vBlock
){
	FLUX_TIMER_RESET;

	Pair<Eigen::Map<const Array1Xs>>
		dCellLR { block.template mapLR<Array1XsCMap,false>(dBlock, er) },
		zCellLR { block.template mapLR<Array1XsCMap,false>(zBlock, er) };

	#ifdef BLOCKWISE_INTEGRATED_EARLY_RETURN_WHEN_DRY
	using S = SolverType;
	if (
		(dCellLR[0] < S::minWaterDepth).all() &&
		(dCellLR[1] < S::minWaterDepth).all()
	){
		return;
	}
	#endif

	EdgesValues ev { reconstruction<recType, orien>(
		er, mesh, block,
		qBlock, dBlock, zBlock, hBlock, vBlock
	) };

	TIME_FLUX(reconstruction);

	Pair<Array3XsMap> flux { hllcScheme<orien>(ev, er.nEdges) };

	TIME_FLUX(hllc);

	slopeSourceToFlux(
		flux,
		ev,
		zCellLR,
		dCellLR,
		rectNormals<orien>(er.nEdges)
	);

	TIME_FLUX(slope);

	registerFlux_vec<recType, orien>(
		fluxes, flux[left], flux[right],
		er, lrBeg_block, block
	);

	TIME_FLUX(registration);
}

template<typename MeshType>
template<ReconstructionType recType, EdgeRange orien, BlockLocation loc>
void SolverThread<MeshType>::registerFlux_vec(
	Array3XsMap& fluxes,
	Ref<Array3Xs> fluxL,
	Ref<Array3Xs> fluxR,
	const EdgeRow<recType,orien>& er,
	const Array2i lrBeg_block,
	const Block<loc>& block
){
	using ER = EdgeRange;
	const scalar& timeStep { solver().time().step() };
	const MeshType& mesh { solver().mesh() };

	constexpr Index nRows {3};

	auto lengths = [&]() -> decltype(auto) {
		if constexpr (std::is_same_v<MeshType, UniMesh>){
			if constexpr ( orien == ER::horizontal )
				return mesh.dx();
			else
				return mesh.dy();
		} else if constexpr (std::is_same_v<MeshType, RectMesh>){
			if constexpr ( orien == ER::horizontal ){
				return mesh.dx()
					.segment(
						block.xBeg() +
						block.template cellOffset<recType, orien>(),
						er.nEdges)
					.transpose()
					.template replicate<nRows,1>();
			} else {
				return mesh.dy()[block.yBeg() + er.i];
			}
		} else {
			return mesh.edges().lengths()
				.segment(er.edgeBeg, er.nEdges)
				.transpose()
				.template replicate<nRows,1>();
		}
	};

	static constexpr int left {0}, right {1};
	Array2i lrBeg_mesh { block.template lrBeg_mesh<recType, orien>(er.i) };

	auto areas = [&]([[maybe_unused]] int side) -> decltype(auto) {
		if constexpr (std::is_same_v<MeshType, UniMesh>){
			return mesh.cells().area();
		} else {
			return mesh.cells().areas()
				.segment( lrBeg_mesh[side], er.nEdges )
				.transpose()
				.template replicate<nRows,1>();
		}
	};

	fluxL *= timeStep * lengths() / areas(left );
	fluxR *= timeStep * lengths() / areas(right);

	fluxes.middleCols(lrBeg_block[left ], er.nEdges) -= fluxL;
	fluxes.middleCols(lrBeg_block[right], er.nEdges) += fluxR;
}


} // namespace hms::swe
